import gradio as gr
from gradio.blocks import Blocks
from gradio.layouts import Row, Column, Accordion, Tabs, TabItem
from gradio.components import (
    ImageEditor, Button, Radio, Slider, State, Markdown, Image
)
from typing import Optional, Tuple

import os
import cv2
import numpy as np
import json

# # Fix configuration
# os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
# os.makedirs("static", exist_ok=True)

from utils.points import (
    point_handler, 
    clear_points_handler
)
from utils.boxes import (
    start_box_handler, 
    update_box_handler, 
    end_box_handler, 
    remove_box_handler, 
    clear_boxes_handler
)
from utils.apis import (
    sam_model_api, 
    inpainting_model_api, 
    enhancement_model_api
)

def create_segmentation_ui():
    """Create segmentation UI with point/box selection using ImageEditor"""
    with Blocks() as ui:
        # Define state variables first
        points_list = State([])
        boxes_list = State([])
        last_point_action = State("add")
        
        with Row():
            with Column():
                input_image = ImageEditor(
                    label="Original Image", 
                    type="numpy",
                    interactive=True,
                    height=600,
                    elem_classes=["scrollable-image"]
                )
                
                # Track initial upload state
                is_initial_upload = State(True)
                
                # Store uploaded images in memory
                uploaded_images = State(None)
                
                def update_original_image(img, is_initial):
                    if img is not None:
                        uploaded_images.value = img.copy()
                        points_list.value = []
                        boxes_list.value = []
                        return img, False
                    return None, is_initial
                    
                input_image.upload(
                    update_original_image,
                    [input_image, is_initial_upload],
                    [uploaded_images, is_initial_upload]
                )
                
                # Handle image edits (point/box placement)
                def handle_image_edit(data):
                    if data is None:
                        return None, None, points_list.value, last_point_action.value
                    
                    # Extract edited image and annotations
                    edited_image = data.get("image")
                    annos = data.get("annotations", [])
                    
                    new_points = []
                    new_boxes = []
                    
                    for anno in annos:
                        if anno.get("type") == "point":
                            x, y = anno.get("x", 0), anno.get("y", 0)
                            new_points.append((x, y, 1))  # Positive point
                        elif anno.get("type") == "rectangle":
                            x, y = anno.get("x", 0), anno.get("y", 0)
                            width, height = anno.get("width", 0), anno.get("height", 0)
                            box = (x, y, x + width, y + height)
                            new_boxes.append(box)
                    
                    # Update state
                    points_list.value = new_points
                    boxes_list.value = new_boxes
                    
                    # Return updated state
                    return edited_image, edited_image, points_list, last_point_action
                
                # Bind the edit handler
                input_image.change(
                    handle_image_edit,
                    inputs=None,
                    outputs=[input_image, uploaded_images, points_list, last_point_action]
                )
                
        with Row():
            # Clear buttons
            clear_points = Button("🧹 Clear Points", min_width=30)
            clear_boxes = Button("🧼 Clear Boxes", min_width=30)
        with Column():
            output_mask = Image(label="Segmentation Mask", interactive=False, height=600)
        
        with Row():
            segment_btn = Button("Segment", variant="primary")
            next_btn = Button("Next", variant="primary")
        
        def clear_points_operation():
            print("Clearing all points")
            points_list.value = []
            # Reset image to original without points
            if uploaded_images.value is not None:
                # Create a copy of the original image
                clean_image = uploaded_images.value.copy()
                # Draw only boxes
                for box in boxes_list.value:
                    if isinstance(box, (list, tuple)) and len(box) == 4:
                        x1, y1, x2, y2 = box
                        cv2.rectangle(clean_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
                return clean_image, points_list.value
            return input_image.value, points_list.value

        clear_points.click(
            fn=clear_points_operation,
            inputs=None,
            outputs=[input_image, points_list]
        )

        # Box operations
        def clear_boxes_operation():
            print("Clearing all boxes")
            boxes_list.value = []
            # Reset image to original without boxes
            if uploaded_images.value is not None:
                # Create a copy of the original image
                clean_image = uploaded_images.value.copy()
                # Draw only points
                for pt in points_list.value:
                    if isinstance(pt, (list, tuple)) and len(pt) >= 2:
                        cv2.drawMarker(
                            clean_image, 
                            (int(pt[0]), int(pt[1])), 
                            (0,0,255), 
                            cv2.MARKER_STAR, 
                            20, 
                            2
                        )
                return clean_image, boxes_list.value
            return input_image.value, boxes_list.value

        clear_boxes.click(
            fn=clear_boxes_operation,
            inputs=None,
            outputs=[input_image, boxes_list]
        )

    return ui, next_btn, output_mask

def create_inpainting_ui():
    """Create inpainting UI for virtual try-on"""
    with Blocks() as ui:
        with Row():
            with Column():
                original_image = Image(label="Original Image", interactive=False)
                segmented_image = Image(label="Segmented Area", interactive=False)
                clothing_image = Image(label="Clothing Image", type="numpy")
            with Column():
                inpainted_image = Image(label="Inpainted Result", interactive=False)
        
        with Row():
            prev_btn = Button("Previous")
            inpaint_btn = Button("Inpaint", variant="primary")
            next_btn = Button("Next", variant="primary")
        
        # Inpaint handler
        def inpaint_handler(original_img, mask_img, clothing_img):
            if None in [original_img, mask_img, clothing_img]:
                return None
            return inpainting_model_api(original_img, mask_img[:,:,0], clothing_img)[0]
        
        # Bind event handlers
        inpaint_btn.click(
            inpaint_handler,
            [original_image, segmented_image, clothing_image],
            inpainted_image
        )
    
    return ui, prev_btn, next_btn, segmented_image

def create_enhancement_ui():
    """Create enhancement UI with options"""
    with Blocks() as ui:
        with Row():
            with Column():
                input_image = Image(label="Input Image", interactive=False)
                with Accordion("Enhancement Settings", open=False):
                    method = Radio(
                        ["detail", "color", "contrast"], 
                        label="Enhancement Method",
                        value="detail"
                    )
                    sigma_s = Slider(1, 20, value=10, label="Detail Strength")
                    sigma_r = Slider(0.01, 0.5, value=0.15, step=0.01, label="Detail Smoothness")
                    contrast = Slider(0.5, 2.0, value=1.2, step=0.1, label="Contrast")
            with Column():
                enhanced_image = Image(label="Enhanced Image", interactive=False)
        
        with Row():
            prev_btn = Button("Previous")
            enhance_btn = Button("Enhance", variant="primary")
            finish_btn = Button("Finish", variant="primary")
        
        # Enhance handler
        def enhance_handler(image, method, sigma_s, sigma极, contrast):
            if image is None:
                return None
            return enhancement_model_api(image, method, sigma_s, sigma_r, contrast)

        # Bind event handlers
        enhance_btn.click(
            enhance_handler,
            [input_image, method, sigma_s, sigma_r, contrast],
            enhanced_image
        )

    return ui, prev_btn

def main():
    """Main application"""
    app = Blocks(title="Virtual Try-On System", css="static/styles.css")
    
    with app:
        Markdown("# Virtual Try-On System")
        
        with Tabs() as tabs:
            with TabItem("1. Segmentation", id="segmentation"):
                seg_ui, seg_next_btn, seg_output_mask = create_segmentation_ui()
            with TabItem("2. Inpainting", id="inpainting"):
                inpaint_ui, inpaint_prev_btn, inpaint_next_btn, inpaint_segmented_image = create_inpainting_ui()
            with TabItem("3. Enhancement", id="enhancement"):
                enhance_ui, enhance_prev_btn = create_enhancement_ui()
        
        tab_id = State("segmentation")
        
        seg_next_btn.click(
            fn=lambda mask: "inpainting",
            inputs=seg_output_mask,
            outputs=tab_id
        ).then(
            lambda tab: gr.update(selected=tab),
            inputs=[tab_id],
            outputs=tabs
        )
        
        inpaint_prev_btn.click(
            fn=lambda: "segmentation",
            inputs=None,
            outputs=tab_id
        ).then(
            lambda tab: gr.update(selected=tab),
            inputs=[tab_id],
            outputs=tabs
        )
        
        inpaint_next_btn.click(
            fn=lambda: "enhancement",
            inputs=None,
            outputs=tab_id
        ).then(
            lambda tab: gr.update(selected=tab),
            inputs=[tab_id],
            outputs=tabs
        )
        
        enhance_prev_btn.click(
            fn=lambda: "inpainting",
            inputs=None,
            outputs=tab_id
        ).then(
            lambda tab: gr.update(selected=tab),
            inputs=[tab_id],
            outputs=tabs
        )
    
    app.launch(
        server_name="127.0.0.1",
        server_port=7860,
        show_error=True,
        favicon_path=None,
        debug=True
    )

if __name__ == "__main__":
    main()
